import os
import pandas as pd
import matplotlib.pyplot as plt
import datetime
from sklearn.metrics import roc_auc_score, classification_report
from my_utils.log import Logger
from sklearn.preprocessing import StandardScaler
import joblib

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15


class BrainDrainModel(object):
    def __init__(self, path):
        logfile_name = 'predict_' + datetime.datetime.now().strftime('%Y%m%d%H%M%S')
        self.logfile = Logger('../', logfile_name).get_logger()
        self.data_source = pd.read_csv(path)


def pred_feature_extract(data, logger):
    # 分析文件
    # print(data.info())
    # print(data.head())
    # 拷贝一份源数据
    feature_data = data.copy()
    # feature_data.drop(['EmployeeNumber', 'Over18', 'StandardHours'], axis=1, inplace=True)
    feature_data = pd.get_dummies(feature_data)
    feature_data['AgeGroup'] = pd.cut(data['Age'], bins=[0, 35, 45, 100], labels=[0, 1, 2]).astype(int)
    feature_data['a']=pd.cut(data['MonthlyIncome'], bins=[0, 3500, 7000, 10000, 15000, 20000], labels=[0, 1, 2, 3, 4]).astype(int)
    feature_data["月收入与工作年限比例"] = data["MonthlyIncome"] / (data["TotalWorkingYears"] + 100)
    feature_data['aa'] = (feature_data['TotalWorkingYears'] + 1) / (feature_data['YearsAtCompany'] + 1)
    feature_data['bb'] = feature_data['YearsSinceLastPromotion'] / (feature_data['YearsAtCompany'] + 1)
    feature_data['cc'] = feature_data['YearsInCurrentRole'] * (feature_data['YearsAtCompany'] + 1)
    feature_data['dd'] = ((feature_data['JobInvolvement'] >= 3) & (feature_data['JobSatisfaction'] <= 2))

    feature_columns = ['Age', 'DistanceFromHome', 'Education',
                       'EnvironmentSatisfaction', 'JobInvolvement', 'JobLevel',
                       'JobSatisfaction', 'MonthlyIncome', 'NumCompaniesWorked',
                       'PercentSalaryHike', 'PerformanceRating', 'RelationshipSatisfaction',
                       'StandardHours', 'StockOptionLevel', 'TotalWorkingYears',
                       'TrainingTimesLastYear', 'WorkLifeBalance', 'YearsAtCompany',
                       'YearsInCurrentRole', 'YearsSinceLastPromotion', 'YearsWithCurrManager',
                       'BusinessTravel_Non-Travel', 'BusinessTravel_Travel_Frequently',
                       'BusinessTravel_Travel_Rarely', 'Department_Human Resources',
                       'Department_Research & Development', 'Department_Sales',
                       'EducationField_Human Resources', 'EducationField_Life Sciences',
                       'EducationField_Marketing', 'EducationField_Medical',
                       'EducationField_Other', 'EducationField_Technical Degree',
                       'Gender_Female', 'Gender_Male', 'JobRole_Healthcare Representative',
                       'JobRole_Human Resources', 'JobRole_Laboratory Technician',
                       'JobRole_Manager', 'JobRole_Manufacturing Director',
                       'JobRole_Research Director', 'JobRole_Research Scientist',
                       'JobRole_Sales Executive', 'JobRole_Sales Representative',
                       'MaritalStatus_Divorced', 'MaritalStatus_Married',
                       'MaritalStatus_Single', 'OverTime_No', 'OverTime_Yes',
                       'AgeGroup', "月收入与工作年限比例", 'aa', 'bb', 'cc', 'dd', 'a']
    return feature_data, feature_columns


# 结果分析画图
def prediction_plot(data):
    pass


def auc_predict(feature_data, feature_name, model, logger):
    x_test = feature_data[feature_name]
    y_test = feature_data['Attrition']
    y_proba = model.predict_proba(x_test)[:, 1]  # 正类概率（AUC必需）
    y_pred = model.predict(x_test)
    auc = roc_auc_score(y_test, y_proba)
    print(auc)
    print(classification_report(y_test, y_pred))
    logger.info(f"{model}模型AUC:{auc}")
    return auc


def auc_predict_l(feature_data, feature_name, model, logger):
    """
    统一的模型预测函数，兼容各种LightGBM模型类型
    参数:
    - feature_data: 包含特征和目标变量的DataFrame
    - feature_name: 特征列名列表
    - model: 训练好的模型（LGBMClassifier或Booster）
    - logger: 日志记录器
    """
    try:
        x_test = feature_data[feature_name]
        y_test = feature_data['Attrition']
        transfer = StandardScaler()
        x_test = transfer.fit_transform(x_test)
        # 根据模型类型选择预测方式
        if hasattr(model, 'predict_proba'):
            # sklearn接口的模型
            y_proba = model.predict_proba(x_test)[:, 1]
        elif hasattr(model, 'predict'):
            # Booster对象
            y_proba = model.predict(x_test)
        else:
            raise ValueError("不支持的模型类型")

        auc = roc_auc_score(y_test, y_proba)

        # 记录日志
        logger.info(f"模型评估 - AUC: {auc:.4f}")
    except Exception as e:
        logger.error(f"预测过程中出错: {str(e)}")
        raise e


if __name__ == '__main__':
    # 加载测试数据
    plp = BrainDrainModel("../data/test.csv")
    feature_data, feature_name = pred_feature_extract(plp.data_source, logger=plp.logfile)

    model = joblib.load('../model/xgboost_20251027.pkl')
    auc_predict(feature_data, feature_name, model, plp.logfile)
